package comlib

import chisel3._
import chisel3.util.log2Up

/** Reduce process class
  *
  * @tparam T Reduce element type
  * @param s Sequence which will be reduce process
  */
class ReduceProcess [T<:Data] (s:Seq[T]) {

  private var slotSize = s.size
  private var lockStage = 0
  private var stage = 1
  private var fillVal:Option[T] = None

  def setSerialize( slotSize:Int, fillVal:Option[T]=None ) = {
    this.slotSize = if(slotSize==0) s.size else slotSize
    this.fillVal = fillVal
    this
  }

  def setBinaryTree( lockStage:Int, stage:Int=1 ) = {
    require( lockStage >= 0 )
    require( (stage>0 && stage<=lockStage) || lockStage==0, s"stage(${stage}) and lockStage(${lockStage}) is not met the reduce condition" )

    this.lockStage = lockStage
    this.stage = stage
    this
  }

    

  /** 串行二叉树 Reduce
    *
    * @param f Reduce function
    */
  def serialBinaryTree( f:(T,T)=>T ) = {

    if( slotSize == s.size ){
      reduce.binaryTree( f, s, lockStage, stage )
    }else{

      // 构建分组
      // val grps = s.groupVec( slotSize, fillVal.get )
      // // 构建TDM序列
      // val (cnt, wrap) = chisel3.util.Counter( true.B, grps.size )
      // // 二叉树比较
      // val ret_cmp = reduce.binaryTree( f, grps(cnt).toSeq, lockStage, stage )
      // // 最终输出
      // reduce.reduceSerial( f, ret_cmp, wrap, true )
    }

  }


  


    /** Get delay value (Unit: clock) of this reduce process
      */
    def delay = {
        // 计算binaryTree延时clock个数
        val ffs = if( lockStage==0 ) 0 else scala.math.ceil((log2Up(s.size) - (lockStage-stage))*1.0 / lockStage).toInt

        if( slotSize == s.size ) ffs      // 直接返回
        else ffs + scala.math.ceil(s.size*1.0/slotSize).toInt + 1 // 叠加串行所需clock个数
    }

}


object reduce{

  def apply[T<:Data](v:Seq[T]): ReduceProcess[T] = new ReduceProcess(v)

  /** Reduce using binary-tree solution
    *
    * @param f Reduce function
    * @param v Reduce operated sequence
    * @param lockStage Recur stage number when locked with a register. Do not insert register when `lockStage`<=0
    * @param stage Initial stage id
    *
    * ==== Example for stage and lockStage ====
    *  - When stage=1 and lockStage=3, the pipeline is: ('|' means wire and '||' means reg)
    *    | Fun-level0 | Fun-level1 || Fun-level2 | Fun-level3 | Fun-level4 || Fun-level5 | Fun-level6 | Fun-level7 || Fun-level8 ...
    *  - When stage=2 and lockStage=3, the pipeline is:
    *    | Fun-level0 || Fun-level1 | Fun-level2 | Fun-level3 || Fun-level4 | Fun-level5 | Fun-level6 || Fun-level7 | Fun-level8 ...
    *  - When stage=3 and lockStage=3, the pipeline is:
    *    || Fun-level0 | Fun-level1 | Fun-level2 || Fun-level3 | Fun-level4 | Fun-level5 || Fun-level6 | Fun-level7 | Fun-level8 ...
    *
    * @note Do not insert register before leaf node
    */
  def binaryTree[T<:Data]( f: (T,T)=>T, v: Seq[T], lockStage: Int=0, stage: Int=1 ):T = v.size match {

      case 1 => v(0)
      case x =>

          // Register locked judgement
          val isLock = (stage == lockStage)
          // Next stage
          val nxtStage = if( isLock ) 1 else stage + 1

          // Recursive
          val left = binaryTree( f, v.take( x/2 ), lockStage, nxtStage )
          val right = binaryTree( f, v.takeRight( x - x/2 ), lockStage, nxtStage )

          if( isLock ) RegNext( f(left, right) ) else f(left, right)
  }

  /** Reduce with a serial of ele and start with `start` signal
    *
    * @param f Reduce function
    * @param ele Reduce element for each clock
    * @param start Reduce start (init) flag
    * @param lock Register output enable flag
    * @return Real-time reduce value
    */
  def reduceSerial[T<:Data]( f:(T,T)=>T, ele:T, start:Bool, lock:Boolean=true ) = {
    val ret_r = Reg(chiselTypeOf(ele))
    val ret_w = Wire(chiselTypeOf(ele))
    when( start ){
      ret_w := ele
    }.otherwise{
      ret_w := f( ret_r, ele )
    }
    ret_r := ret_w
    if( lock ) ret_r else ret_w
  }

}


